package com.nanohadoop.yarn;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 基于轮盘赌算法的资源调度器
 * 根据应用权重随机选择应用进行资源分配
 */
public class WeightedResourceScheduler implements ResourceScheduler {
    private static final Logger log = LoggerFactory.getLogger(WeightedResourceScheduler.class);

    private final Map<String, NodeManager> nodeManagers;
    private final List<ApplicationMaster> applicationQueue;
    private final Map<String, ApplicationMaster> applications;
    private final Map<String, Integer> applicationWeights;
    private int nextContainerId;
    private int nextApplicationId;
    private final ScheduledExecutorService schedulerService;
    private volatile boolean running;
    private final Random random;

    public WeightedResourceScheduler() {
        this.nodeManagers = new ConcurrentHashMap<>();
        this.applicationQueue = new ArrayList<>();
        this.applications = new ConcurrentHashMap<>();
        this.applicationWeights = new ConcurrentHashMap<>();
        this.nextContainerId = 0;
        this.nextApplicationId = 0;
        this.schedulerService = Executors.newSingleThreadScheduledExecutor();
        this.running = false;
        this.random = new Random();
        log.debug("WeightedResourceScheduler initialized");
    }

    @Override
    public void start() {
        if (running) {
            return;
        }
        running = true;
        // 每100毫秒执行一次调度
        schedulerService.scheduleAtFixedRate(this::schedule,
                0, 
                ResourceSchedulerConfig.SCHEDULE_INTERVAL_MS,
                TimeUnit.MILLISECONDS);
        log.debug("WeightedResourceScheduler started");
    }

    @Override
    public void stop() {
        if (!running) {
            return;
        }
        running = false;
        schedulerService.shutdown();
        try {
            if (!schedulerService.awaitTermination(5, TimeUnit.SECONDS)) {
                schedulerService.shutdownNow();
            }
        } catch (InterruptedException e) {
            schedulerService.shutdownNow();
        }
        log.debug("WeightedResourceScheduler stopped");
    }

    @Override
    public void registerNodeManager(NodeManager nodeManager) {
        nodeManagers.put(nodeManager.getNodeId(), nodeManager);
        log.debug("NodeManager registered: " + nodeManager.getNodeId());
    }

    @Override
    public String submitApplication(String name, String user, Resource requiredResource, int numContainers) {
        String applicationId = "app_" + nextApplicationId++;
        ApplicationMaster application = new ApplicationMaster(applicationId, name, user, requiredResource, numContainers);
        applications.put(applicationId, application);
        applicationQueue.add(application);
        // 默认权重设为numContainers，也可根据实际需求调整
        applicationWeights.put(applicationId, numContainers);
        log.debug("Application submitted: " + applicationId + ", name: " + name + ", user: " + user + ", required resources per container: " + requiredResource.getMemoryMB() + "MB/" + requiredResource.getCpuCores() + "cores, total containers: " + numContainers + ", weight: " + numContainers);
        return applicationId;
    }

    @Override
    public void schedule() {
        if (!running) {
            log.debug("Scheduler is not running, skipping schedule");
            return;
        }

        if (applicationQueue.isEmpty()) {
            log.debug("No applications in queue, skipping schedule");
            return;
        }

        if (nodeManagers.isEmpty()) {
            log.debug("No node managers available, skipping schedule");
            return;
        }

        // 轮盘赌算法选择应用
        ApplicationMaster selectedApp = selectApplicationByRouletteWheel();
        if (selectedApp == null) {
            log.debug("No application selected by roulette wheel");
            return;
        }

        log.debug("Scheduling application: " + selectedApp.getApplicationId() + ", current status: " + selectedApp.getStatus());

        if (selectedApp.getStatus() != ApplicationMaster.ApplicationStatus.SUBMITTED &&
                selectedApp.getStatus() != ApplicationMaster.ApplicationStatus.RUNNING) {
            log.debug("Application " + selectedApp.getApplicationId() + " has status " + selectedApp.getStatus() + ", removing from queue");
            applicationQueue.remove(selectedApp);
            applicationWeights.remove(selectedApp.getApplicationId());
            return;
        }

        // 检查是否还需要分配容器
        int allocatedContainers = selectedApp.getAllocatedContainers().size();
        int numContainers = selectedApp.getNumContainers();
        log.debug("Application " + selectedApp.getApplicationId() + ": allocated " + allocatedContainers + "/" + numContainers + " containers");

        if (allocatedContainers >= numContainers) {
            log.debug("All containers allocated for application " + selectedApp.getApplicationId() + ", setting status to RUNNING");
            selectedApp.setStatus(ApplicationMaster.ApplicationStatus.RUNNING);
            applicationQueue.remove(selectedApp);
            applicationWeights.remove(selectedApp.getApplicationId());
            return;
        }

        Resource requiredResource = selectedApp.getRequiredResource();
        boolean foundNode = false;
        for (NodeManager nodeManager : nodeManagers.values()) {
            log.debug("Checking node " + nodeManager.getNodeId() + " for resources");
            if (nodeManager.hasEnoughResource(requiredResource)) {
                foundNode = true;
                log.debug("Node " + nodeManager.getNodeId() + " has enough resources");
                // 创建容器
                String containerId = "container_" + nextContainerId++;
                Container container = new Container(
                        containerId,
                        nodeManager.getNodeId(),
                        requiredResource,
                        selectedApp.getApplicationId()
                );

                // 分配容器
                log.debug("Attempting to allocate container " + containerId + " on node " + nodeManager.getNodeId());
                boolean allocationResult = nodeManager.allocateContainer(container);
                log.debug("Container allocation result: " + allocationResult);

                if (allocationResult) {
                    selectedApp.addContainer(container);
                    log.debug("Container " + containerId + " allocated to application " + selectedApp.getApplicationId());
                    selectedApp.setStatus(ApplicationMaster.ApplicationStatus.RUNNING);
                    log.debug("Application " + selectedApp.getApplicationId() + " status changed to RUNNING");

                    // 检查是否所有容器都已分配
                    allocatedContainers = selectedApp.getAllocatedContainers().size();
                    if (allocatedContainers >= numContainers) {
                        log.debug("All containers allocated for application " + selectedApp.getApplicationId() + ", removing from queue");
                        applicationQueue.remove(selectedApp);
                        applicationWeights.remove(selectedApp.getApplicationId());
                    }
                } else {
                    log.debug("Failed to allocate container " + containerId + " on node " + nodeManager.getNodeId());
                    log.debug("Node status: running=" + nodeManager.isRunning());
                    log.debug("Node available resources: " + nodeManager.getAvailableResource().getMemoryMB() + "MB memory, " + nodeManager.getAvailableResource().getCpuCores() + " CPU cores");
                    log.debug("Node required resources: " + requiredResource.getMemoryMB() + "MB memory, " + requiredResource.getCpuCores() + " CPU cores");
                }
                // 只分配一个容器，然后退出循环
                break;
            }
        }

        if (!foundNode) {
            log.debug("No node with enough resources found for application " + selectedApp.getApplicationId());
        }

        // 如果应用状态还是SUBMITTED，尝试再次设置为RUNNING
        if (selectedApp.getStatus() == ApplicationMaster.ApplicationStatus.SUBMITTED) {
            log.debug("Application " + selectedApp.getApplicationId() + " still SUBMITTED after scheduling attempt");
        }
    }

    /**
     * 轮盘赌算法选择应用
     */
    private ApplicationMaster selectApplicationByRouletteWheel() {
        // 过滤掉非SUBMITTED和非RUNNING状态的应用
        List<ApplicationMaster> activeApps = new ArrayList<>();
        for (ApplicationMaster app : applicationQueue) {
            if (app.getStatus() == ApplicationMaster.ApplicationStatus.SUBMITTED ||
                    app.getStatus() == ApplicationMaster.ApplicationStatus.RUNNING) {
                activeApps.add(app);
            }
        }

        if (activeApps.isEmpty()) {
            return null;
        }

        // 计算权重总和
        int totalWeight = 0;
        for (ApplicationMaster app : activeApps) {
            totalWeight += applicationWeights.getOrDefault(app.getApplicationId(), 1);
        }

        if (totalWeight <= 0) {
            return activeApps.get(random.nextInt(activeApps.size()));
        }

        // 生成随机数
        int randomValue = random.nextInt(totalWeight);

        // 轮盘赌选择
        int currentSum = 0;
        for (ApplicationMaster app : activeApps) {
            int weight = applicationWeights.getOrDefault(app.getApplicationId(), 1);
            currentSum += weight;
            if (randomValue < currentSum) {
                return app;
            }
        }

        // 如果出现意外情况，返回随机应用
        return activeApps.get(random.nextInt(activeApps.size()));
    }

    @Override
    public void completeApplication(String applicationId) {
        ApplicationMaster application = applications.get(applicationId);
        if (application != null) {
            log.debug("Completing application: " + applicationId + ", current status: " + application.getStatus());
            application.setStatus(ApplicationMaster.ApplicationStatus.COMPLETED);
            // 释放所有容器资源
            for (Container container : application.getAllocatedContainers()) {
                NodeManager nodeManager = nodeManagers.get(container.getNodeId());
                if (nodeManager != null) {
                    nodeManager.releaseContainer(container.getContainerId());
                    log.debug("Released container: " + container.getContainerId() + " for application: " + applicationId);
                }
            }
            applicationQueue.remove(application);
            applicationWeights.remove(applicationId);
            log.debug("Application completed: " + applicationId);
        } else {
            log.debug("Application not found: " + applicationId);
        }
    }

    @Override
    public ApplicationMaster getApplication(String applicationId) {
        return applications.get(applicationId);
    }

    @Override
    public Map<String, NodeManager> getNodeManagers() {
        return nodeManagers;
    }

    @Override
    public Map<String, ApplicationMaster> getApplications() {
        return applications;
    }

    /**
     * 设置应用权重
     */
    public void setApplicationWeight(String applicationId, int weight) {
        if (applications.containsKey(applicationId)) {
            applicationWeights.put(applicationId, Math.max(1, weight)); // 权重至少为1
            log.debug("Set weight for application " + applicationId + ": " + weight);
        }
    }
}

// 添加配置类
class ResourceSchedulerConfig {
    public static final int SCHEDULE_INTERVAL_MS = 100;
}